import gradio as gr
import numpy as np
import pickle
from PIL import Image
from skimage.transform import resize

# 加载保存的KNN模型
with open("best_knn_model.pkl", "rb") as f:
    knn = pickle.load(f)

def preprocess(img):
    """
    优化后的预处理函数，通过裁剪和居中填充提升识别准确率。
    """
    # 1. 提取图像数据
    # 新版 Gradio 的 Sketchpad 返回 dict，手写内容在 'layers'
    if isinstance(img, dict) and "layers" in img and img["layers"]:
        img = img["layers"][0]
    elif isinstance(img, dict) and "image" in img:
        img = img["image"]

    if img is None:
        return None

    # 2. 转换为灰度 NumPy 数组
    if isinstance(img, Image.Image):
        arr = np.array(img.convert("L"), dtype=np.float32)
    else:
        arr = np.array(img, dtype=np.float32)

    # 如果是 RGBA 图像，手写内容在 Alpha 通道
    if arr.ndim == 3 and arr.shape[2] == 4:
        arr = arr[:, :, 3]
    elif arr.ndim == 3:
        arr = arr[..., 0]

    # 3. 裁剪空白边缘
    # 找到有笔迹的行和列
    rows = np.any(arr, axis=1)
    cols = np.any(arr, axis=0)
    if not np.any(rows) or not np.any(cols):
        return None # 图像为空
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    
    # 裁剪
    arr = arr[rmin:rmax+1, cmin:cmax+1]

    # 4. 填充为正方形并居中
    h, w = arr.shape
    if h != w:
        # 计算需要填充的大小
        diff = abs(h - w)
        pad_before = diff // 2
        pad_after = diff - pad_before
        # 在短边上填充
        if h > w:
            padding = ((0, 0), (pad_before, pad_after))
        else:
            padding = ((pad_before, pad_after), (0, 0))
        arr = np.pad(arr, padding, mode='constant', constant_values=0)

    # 5. 归一化、缩放和格式化
    # 归一化到 0-1
    if arr.max() > 1:
        arr = arr / 255.0

    # 缩放到 8x8
    img_resized = resize(arr, (8, 8), anti_aliasing=True)
    
    # 像素值映射到 0-16
    img_rescaled = (img_resized * 16).astype(np.float64)
    
    return img_rescaled.reshape(1, -1)

# 预测函数
def predict_digit(img):
    arr = preprocess(img)
    if arr is None:
        return "无效图片或空白输入"
    pred = knn.predict(arr)[0]
    return str(pred)

# 使用 gr.Sketchpad 作为手写板输入
with gr.Blocks() as demo:
    gr.Markdown("# 手写数字识别 (KNN)")
    with gr.Row():
        with gr.Column():
            canvas = gr.Sketchpad(label="请在此手写数字")
            btn = gr.Button("识别")
        with gr.Column():
            output = gr.Textbox(label="预测结果")
    btn.click(fn=predict_digit, inputs=canvas, outputs=output)

print("Gradio App is running...")
demo.launch()
